import pandas as pd
import numpy as np
from sqlbatis import SQLBatis


local_database = 'postgresql://postgres:9755@127.0.0.1/invest'
db = SQLBatis(local_database)





sql_get_close_and_mcap = """
select
    fic_code,close,mktcap
from(
select
    fic_code
    ,rank() over(partition by fic_code order by date desc) as day_rank
    ,close
    ,mktcap
from equity_transaction_daily
where date<= :date
    and fic_code in :code) t1
where day_rank=1;
"""
@db.query(sql_get_close_and_mcap)
def close_and_mcap(date,code):
    """
    获得指定日期指定股票的收盘价和当日流通市值，如果当日某支股票没有数据，则用前一有记录的日期的数据来代替
    :param date:
    :param code:
    :return: 返回fic_code ,close,mktcap
    """
    pass


sql_get_daily_r="""
select fic_code,date,(close-lsd_close)/lsd_close as r
from (
select stock_pool.fic_code,date,
       close,lag(close,1) over(partition by t1.fic_code order by date asc) as lsd_close
from stock_pool
left join equity_transaction_daily as t1
on stock_pool.fic_code=t1.fic_code
where date between '2011-12-20' and '2013-01-04') t1
where date between '2011-12-21' and '2013-01-04'
and fic_code in :fic_code;
"""
@db.query(sql_get_daily_r)
def get_daily_r(fic_code):
    """
    获得2011-12-21至2013-01-04(即建仓前一年)初始股票池中每支股票的每日简单收益率(当日收盘价-昨日收盘价）/昨日收盘价
    :param fic_code:
    :return:fic_code,date,r
    """
    pass



sql_avg_simple_r = """
select fic_code,avg(simple_r) as r
from(
select fic_code,value_add/
lag(security_mkt_value,1) over(partition by fic_code order by date) as simple_r
from port_cons_daily
where fic_code in :fic_code
and date between :start_date and :end_date) t1 group by fic_code

"""
@db.query(sql_avg_simple_r)
def avg_simple_r(fic_code,start_date,end_date):
    """
    返回一段时间内，组合中每支股票的持有期平均收益率
    :param fic_code:
    :param start_date:
    :param end_date:
    :return: 每支股票在该时间段内的日平均收益率
    """
    pass

sql_shares_valueadd_mkv = """
select fic_code,date,shares,security_mkt_value as mkv,value_add
from port_cons_daily
where date between :start_date and :end_date
"""
@db.query(sql_shares_valueadd_mkv)
def shares_valueadd_mkv(start_date,end_date):
    """
    指定时间段内每支股票每天的shares、value_add、mkv
    :param start_date:
    :param end_date:
    :return:
    """
    pass


sql_market_date = """
select distinct date
from equity_transaction_daily
where date between :start_date and :end_date
order by date
"""
@db.query(sql_market_date)
def get_market_date(start_date,end_date):
    """
    返回指定时间段内的交易日日期
    :param start_date:
    :param end_date:
    :return:
    """
    pass

sql_dump = "insert into port_cons_daily" \
           "(portfolio_id,date, fic_code ,security_mkt_value, shares,weight,cash_flow,value_add) " \
           "values(:portfolio_id,:date, :fic_code,:security_mkt_value, :shares, :weight, :cash_flow, :value_add);"
@db.bulk_query(sql_dump)
def dump_to_db(a):
    pass

sql_dump_cash = """
insert into port_cash(date,cash)
values(:date,:cash)
"""
@db.query(sql_dump_cash)
def dump_cash(date,cash):
    pass






def buy_stock(r_close,purchase):
    """
    :param r_close:dataframe(fic_code,r,close) r为持有期收益率
    :param purchase: 有多少钱用于购买stock，非负数
    :return: dataframe(fic_code,shares)每支股票需要增加的shares cash调仓之后组合还剩的现金
    """
    print('加仓：{}元'.format(purchase))
    data = r_close.copy()
    try:
        r_scale= (r_close['r']-r_close['r'].min())/(r_close['r'].max()-r_close['r'].min())
        r_scale[len(r_scale) - 1] = r_scale[len(r_scale) - 2]
        data['weight'] = r_scale/r_scale.sum()
    except:
        data['weight'] = 1/len(data)
    price = np.sum(data['weight']*data['close']*100) # 1份的单价 一份包含0.2手stockA 0.3手stockB 0.5手stockC
    cnt = np.floor(purchase/price) #常数，一共能买多少份
    try:
        cnt = float(cnt)
    except:
        cnt = cnt[0]

    print('一共能买{}份'.format(cnt))

    data['shares_change'] = [np.floor(cnt*i)*100 for i in data['weight']]
    return data[['fic_code','shares_change']]

def drop_stock(r_close,shares,withdraw,mode,stock_need_drop):
    """

    :param r_close: dataframe 至少包含fic_code,到目前为止每支股票的持有期收益(日平均)、每支股票当日close
    :param shares:目前每支股票的shares
    :param cash:现在组合还剩多少现金
    :param withdraw: 当赎回日时，withdraw为指定赎回金额；当为主动调仓日时，withdraw为自定金额
    :param mode: mode=1为不满足10%要求减仓方式
    :param stock_need_drop:mode为1时需要减仓的股票code
    :return:
    """
    data = r_close.copy()
    if mode == 1: # 不满足那两个10%要求
        data['shares_change'] = 0
        if len(stock_need_drop)>0:
            change = [i in stock_need_drop for i in list(data.fic_code)]
            data.loc[change,'shares_change'] = 50000


    else: # 赎回或者主动调仓
        data = data.merge(shares,on='fic_code')
        data.sort_values('r',inplace=True,ascending=True) # data按收益率升序排列
        data.index = list(range(len(data)))
        new_shares = [min(1000,i) for i in data['shares']] # new_shares和data顺序一样
        shares_change = data['shares']-new_shares # shares_change和data顺序一样

        drop_amt = shares_change*data['close']

        drop_ac_amt = np.cumsum(drop_amt)
        for i in range(len(drop_ac_amt)):
            if drop_ac_amt[i]>=withdraw:
                stop_idx = i
                break
        change = np.array([1]*(stop_idx+1)+[0]*(len(drop_ac_amt)-(stop_idx+1)))

        data['shares_change'] = shares_change*change


    return data[['fic_code','shares_change']]

def crt1(share_close_mkv,cash):
    """
    判断组合中股票的市值是否小于组合总价值的80%
    :param share_close_mkv: 当前的share和close,以及每支股票在市场中的总值
    :param cash: 当前的现金
    :return: 如果小于80% 返回True 否则返回False
    """
    return np.sum(share_close_mkv['close']*share_close_mkv['shares'])/cash<4

def crt2(share_close_mkv):
    """
    个股持仓市值是否超过该⽀股票总流通市值10%
    :param share_close_mkv: 当前的share和close,以及每支股票在市场中的总值
    :return:
    """
    mkv = share_close_mkv['close'] * share_close_mkv['shares']
    stock_need_drop = share_close_mkv.loc[np.where(mkv > share_close_mkv['mktcap'] * 0.1),'fic_code'].values
    # print('不满足crt2,需要减仓的股票是:{}'.format(stock_need_drop))

    return stock_need_drop

def crt3(share_close_mkv):
    """
    个股持仓市值是否超过投资组合总市值10%
    :param share_close_mkv: 当前的share和close,以及每支股票在市场中的总值
    :return:
    """
    mkv = share_close_mkv['close']*share_close_mkv['shares']
    stock_need_drop = share_close_mkv.loc[np.where(mkv > mkv.sum()* 0.1),'fic_code'].values
    # print('不满足crt3,需要减仓的股票是:{}'.format(stock_need_drop))
    return stock_need_drop

